import numpy as np
import matplotlib.pyplot as plt
import os
import json

# 读取 config.json 文件中的内容
config_file = 'input/config.json'

# 检查 config.json 文件是否存在
if os.path.exists(config_file):
    with open(config_file, 'r') as f:
        config_data = json.load(f)
        n_values = config_data.get("n", [])
        boundary_type = config_data.get("boundary_conditions", {}).get("type", "Unknown")  # 获取边界条件类型
else:
    raise FileNotFoundError(f"{config_file} not found.")

# 确保保存图像的目录存在
os.makedirs('figure', exist_ok=True)

# 遍历每个 n 值
for n in n_values:
    # 构造 CSV 文件名
    filename = f'output/solution_{n}.csv'

    # 检查文件是否存在
    if os.path.exists(filename):
        # 读取 CSV 文件中的数据
        data = np.loadtxt(filename, delimiter=',')

        # 获取网格大小，假设解的数组形状为 (n+2) x (n+2)，因为包括边界
        grid_size = n + 2  # 包括边界条件
        Z = data.reshape((grid_size, grid_size))

        # 生成 X 和 Y 坐标网格
        x = np.linspace(0, 1, grid_size)
        y = np.linspace(0, 1, grid_size)
        X, Y = np.meshgrid(x, y)

        # 绘制热图
        plt.figure(figsize=(8, 6))
        plt.contourf(X, Y, Z, cmap='viridis')
        plt.colorbar(label="Function Value")
        plt.title(f"Solution for n={n}, Boundary Type={boundary_type}")
        plt.xlabel('x')
        plt.ylabel('y')

        # 根据边界条件类型决定后缀
        suffix = boundary_type.lower()  # 使用小写形式作为后缀

        # 保存图像，加入边界条件类型的后缀
        plt.savefig(f'figure/solution_{n}_{suffix}.png')
        plt.close()
    else:
        print(f"Warning: {filename} not found.")
